{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Getting Started with a simple conditional distribution model -- Python"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "colab_button",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop//Users/jiazhenghao/CodingProjects/research/gtsam/python/gtsam/examples/DiscreteMotionModel.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "license_cell",
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "source": [
        "GTSAM Copyright 2010-2022, Georgia Tech Research Corporation,\n",
        "Atlanta, Georgia 30332-0415\n",
        "All Rights Reserved\n",
        "\n",
        "Authors: Frank Dellaert, et al. (see THANKS for the full author list)\n",
        "\n",
        "See LICENSE for the license information"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 57,
      "id": "colab_import",
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    import google.colab\n",
        "    %pip install --quiet gtsam-develop\n",
        "except ImportError:\n",
        "    pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*For further information on theories and implementations, refer to the robotics book: [https://www.roboticsbook.org/](https://www.roboticsbook.org/)*\n",
        "\n",
        "\n",
        "In this notebook, we walk through an simple example of: \n",
        "1. Representing discrete random variables in gtsam\n",
        "2. Building a bayes net with the discrete random variables\n",
        "3. Visualizing the bayes net we built\n",
        "4. Doing effecient inference with gtsam FactorGraphs (reducing the inference time from exponential to polynomial)!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Other Installs and Imports\n",
        "\n",
        "Some of the visualizations below use [graphviz](https://graphviz.org/). That is not super easy to install except when using conda, in which case it is a one-liner:\n",
        "```bash\n",
        "conda install -c conda-forge python-graphviz\n",
        "```\n",
        "It is also installed by default on Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 58,
      "metadata": {},
      "outputs": [],
      "source": [
        "try:\n",
        "    import google.colab\n",
        "    %pip install -U -q gtbook # this installs more visualization mojo, but needs graphviz!\n",
        "except ImportError:\n",
        "    pass"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 59,
      "metadata": {},
      "outputs": [],
      "source": [
        "import gtsam\n",
        "import gtbook\n",
        "from gtbook import display"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Representing Discrete Random Variables"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "For this simple example, think of a Robot that is moving between different rooms for several seconds. After 1, 2, 3, ... n seconds, the robot wouldn't be sure where it is any more. So for each timestamp t (1,2,3, ...n), we use a random variable $X_t$ to represent the probability of the robot in each room.\n",
        "\n",
        "We have only 5 rooms `['Living Room', 'Kitchen', 'Office', 'Hallway', 'Dining Room']` variable X can take on. So each $X_t$ is a probability distribution over these 5 states. Let's represent that in gtsam. \n",
        "\n",
        "\n",
        "<figure id=\"fig:House\">\n",
        "<img src=\"https://github.com/gtbook/robotics/blob/main/Figures3/S31-House.png?raw=1\" alt=\"\" />\n",
        "<figcaption>The floor plan of a house in which our hypothetical vacuum cleaning robot will operate.</figcaption>\n",
        "</figure>\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 60,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(0, 5)\n",
            "(3, 5)\n"
          ]
        }
      ],
      "source": [
        "# We want to create random variables X1, X2, X3, ...\n",
        "from gtbook.discrete import Variables\n",
        "VARIABLES = Variables()\n",
        "\n",
        "# The possible states of the random variable\n",
        "room_states = ['Living Room', 'Kitchen', 'Office', 'Hallway', 'Dining Room']\n",
        "\n",
        "# create a discrete variable with states\n",
        "\n",
        "#                   name,    states\n",
        "X1 = VARIABLES.discrete('X1', room_states)\n",
        "X2 = VARIABLES.discrete('X2', room_states)\n",
        "X3 = VARIABLES.discrete('X3', room_states)\n",
        "X4 = VARIABLES.discrete('X4', room_states)\n",
        "\n",
        "print(X1)\n",
        "print(X4)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We have just created random variables $X1, X2, X3, X4$. Each one is a discrete variable that has probability of being in each of the room_states. \n",
        "\n",
        "The printed output we see for X1 and X2 is a tuple (key, cardinality). The key is the gtsam internal representation of the variable. Cardinality in our case is just the number of rooms: 5. \n",
        "\n",
        "At any time, if you are confused about how to use a function or how to interpret the result, you can run"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 61,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Help on method discrete in module gtbook.discrete:\n",
            "\n",
            "discrete(name: str, domain: list[str]) -> tuple[int, int] method of gtbook.discrete.Variables instance\n",
            "    Create a variable with given name and discrete domain of named values.\n",
            "\n",
            "    Args:\n",
            "        name (str): name of the variable.\n",
            "        domain (list[str]): names for the different values.\n",
            "\n",
            "    Returns:\n",
            "        DiscreteKey, i.e., (gtsam.Key, cardinality)\n",
            "\n"
          ]
        }
      ],
      "source": [
        "help(VARIABLES.discrete)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Discrete Series\n",
        "In the previous example, it's a bit tedious to create each variable one by one. We can use `gtsam.Variables.discrete_series` to create a series of discrete variables with the same cardinality."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 62,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{1: (6341068275337658369, 5), 2: (6341068275337658370, 5), 3: (6341068275337658371, 5), 4: (6341068275337658372, 5)}\n",
            "X1: (6341068275337658369, 5)\n",
            "X4: (6341068275337658372, 5)\n"
          ]
        }
      ],
      "source": [
        "X = VARIABLES.discrete_series('X', range(1, 5), room_states)\n",
        "print(X)\n",
        "print(f\"X1: {X[1]}\")\n",
        "print(f\"X4: {X[4]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Here, we see the keys got much longer, but the object is really exactly the same as what we had above. For more details on how this value is constructed, see [here](https://www.roboticsbook.org/S31_vacuum_state.html#gtsam-101). \n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Assigning Probabilities\n",
        "We can now assign probabilities to each variable. Let's initialize X1 to be in the Living Room with 80% probability, and the rest of the rooms with 5% probability each.\n",
        "\n",
        "1. Create a Graph `gtsam.DiscreteFactorGraph` to record the assignments. \n",
        "2. Add the variables to the graph with the assigned probabilities.\n",
        "\n",
        "*Don't get frustrated by the term `FactorGraph`, it's just a graph that holds variables and their probabilities. To learn more about FactorGraphs, and how we use it to perform optimization, view the [roboticsbook](https://www.roboticsbook.org/S34_vacuum_perception.html#converting-bayes-nets-into-factor-graphs) and the [gtsam documentation](https://borglab.github.io/gtsam/discretefactorgraph/).*\n",
        "\n",
        "<!-- TODO: Are there better ways to remove variables?  -->"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 63,
      "metadata": {},
      "outputs": [],
      "source": [
        "# create graph\n",
        "graph = gtsam.DiscreteFactorGraph()\n",
        "\n",
        "# adding probabilities for X1\n",
        "# X = ['Living Room', 'Kitchen', 'Office', 'Hallway', 'Dining Room']\n",
        "prob = [0.8, 0.05, 0.05, 0.05, 0.05]  # P(X1)\n",
        "graph.add(X[1], prob)  # P(X1)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 64,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div><p><tt>DiscreteFactorGraph</tt> of size 1</p><p>factor 0:</p><div>\n",
              "<table class='DecisionTreeFactor'>\n",
              "  <thead>\n",
              "    <tr><th>X1</th><th>value</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>Living Room</th><td>0.8</td></tr>\n",
              "    <tr><th>Kitchen</th><td>0.05</td></tr>\n",
              "    <tr><th>Office</th><td>0.05</td></tr>\n",
              "    <tr><th>Hallway</th><td>0.05</td></tr>\n",
              "    <tr><th>Dining Room</th><td>0.05</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n"
            ],
            "text/plain": [
              "<gtbook.display.pretty at 0x10ce2a960>"
            ]
          },
          "execution_count": 64,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Let's view the assignments by printing the graph: \n",
        "\n",
        "def pretty(obj):\n",
        "    return gtbook.display.pretty(obj, VARIABLES) # function to display \n",
        "\n",
        "pretty(graph)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 65,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 12.0.0 (0)\n",
              " -->\n",
              "<!-- Pages: 1 -->\n",
              "<svg width=\"62pt\" height=\"84pt\"\n",
              " viewBox=\"0.00 0.00 62.00 83.60\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 79.6)\">\n",
              "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-79.6 58,-79.6 58,4 -4,4\"/>\n",
              "<!-- var6341068275337658369 -->\n",
              "<g id=\"node1\" class=\"node\">\n",
              "<title>var6341068275337658369</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-57.6\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-52.55\" font-family=\"Times,serif\" font-size=\"14.00\">X1</text>\n",
              "</g>\n",
              "<!-- factor0 -->\n",
              "<g id=\"node2\" class=\"node\">\n",
              "<title>factor0</title>\n",
              "<ellipse fill=\"black\" stroke=\"black\" cx=\"27\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
              "</g>\n",
              "<!-- var6341068275337658369&#45;&#45;factor0 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>var6341068275337658369&#45;&#45;factor0</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M27,-39.28C27,-25.77 27,-8.54 27,-3.96\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<gtbook.display.show at 0x10ce29ac0>"
            ]
          },
          "execution_count": 65,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# This cell won't run if you don't have graphviz installed!\n",
        "display.show(graph)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 66,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div><p><tt>DiscreteFactorGraph</tt> of size 1</p><p>factor 0:</p><div>\n",
              "<table class='DecisionTreeFactor'>\n",
              "  <thead>\n",
              "    <tr><th>X1</th><th>value</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>Living Room</th><td>80</td></tr>\n",
              "    <tr><th>Kitchen</th><td>5</td></tr>\n",
              "    <tr><th>Office</th><td>5</td></tr>\n",
              "    <tr><th>Hallway</th><td>5</td></tr>\n",
              "    <tr><th>Dining Room</th><td>5</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n"
            ],
            "text/plain": [
              "<gtbook.display.pretty at 0x10ce3da30>"
            ]
          },
          "execution_count": 66,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Alternatively, you can also specify the probabilities in string format!\n",
        "# Plus, they don't have to sum to 1, gtsam will normalize them for you. So you can also give them in terms of frequency\n",
        "\n",
        "graph2 = gtsam.DiscreteFactorGraph()\n",
        "# This works:\n",
        "# graph2.add(X[1], \"0.8 0.05 0.05 0.05 0.05\")  # P(X1)\n",
        "\n",
        "# This also works:\n",
        "graph2.add(X[1], \"80 5 5 5 5\")  # P(X1) in terms of frequency\n",
        "\n",
        "pretty(graph2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Conditional Probability\n",
        "\n",
        "Let's now move between the rooms! Let's allow the robot to move *up, down, left, right*. And after taking each action, the room X the robot is in change. \n",
        "To do this, introduce a new variable $A$ that represent what actions we took, and make next next room the robot with be in $X_{t+1}$, depend (conditional) on the action it takes $A_t$, and where it currently is $X_t$. \n",
        "\n",
        "Here we focus on building the **conditional probability**. To learn more about how this simulated robot environment works, visit [robotics-book-ch3](https://www.roboticsbook.org/S30_vacuum_intro.html)\n",
        "\n",
        "Below is the probability table we want to model. Given, current room X1, action A1, what's the probability of being in each room in X2.\n",
        "<figure id=\"fig:House\">\n",
        "<img src=\"https://raw.githubusercontent.com/borglab/gtsam/refs/heads/develop/doc/images/state_action_cond_prob.png \" alt=\"\" />\n",
        "<figcaption></figcaption>\n",
        "</figure>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 67,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div><p><tt>DiscreteFactorGraph</tt> of size 1</p><p>factor 0:</p><div>\n",
              "<p>  <i>P(X2|X1,A1):</i></p>\n",
              "<table class='DiscreteConditional'>\n",
              "  <thead>\n",
              "    <tr><th><i>X1</i></th><th><i>A1</i></th><th>Living Room</th><th>Kitchen</th><th>Office</th><th>Hallway</th><th>Dining Room</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>Living Room</th><th>L</th><td>1</td><td>0</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Living Room</th><th>R</th><td>0.2</td><td>0.8</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Living Room</th><th>U</th><td>1</td><td>0</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Living Room</th><th>D</th><td>0.2</td><td>0</td><td>0</td><td>0.8</td><td>0</td></tr>\n",
              "    <tr><th>Kitchen</th><th>L</th><td>0.8</td><td>0.2</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Kitchen</th><th>R</th><td>0</td><td>1</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Kitchen</th><th>U</th><td>0</td><td>1</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Kitchen</th><th>D</th><td>0</td><td>0.2</td><td>0</td><td>0</td><td>0.8</td></tr>\n",
              "    <tr><th>Office</th><th>L</th><td>0</td><td>0</td><td>1</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Office</th><th>R</th><td>0</td><td>0</td><td>0.2</td><td>0.8</td><td>0</td></tr>\n",
              "    <tr><th>Office</th><th>U</th><td>0</td><td>0</td><td>1</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Office</th><th>D</th><td>0</td><td>0</td><td>1</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Hallway</th><th>L</th><td>0</td><td>0</td><td>0.8</td><td>0.2</td><td>0</td></tr>\n",
              "    <tr><th>Hallway</th><th>R</th><td>0</td><td>0</td><td>0</td><td>0.2</td><td>0.8</td></tr>\n",
              "    <tr><th>Hallway</th><th>U</th><td>0.8</td><td>0</td><td>0</td><td>0.2</td><td>0</td></tr>\n",
              "    <tr><th>Hallway</th><th>D</th><td>0</td><td>0</td><td>0</td><td>1</td><td>0</td></tr>\n",
              "    <tr><th>Dining Room</th><th>L</th><td>0</td><td>0</td><td>0</td><td>0.8</td><td>0.2</td></tr>\n",
              "    <tr><th>Dining Room</th><th>R</th><td>0</td><td>0</td><td>0</td><td>0</td><td>1</td></tr>\n",
              "    <tr><th>Dining Room</th><th>U</th><td>0</td><td>0.8</td><td>0</td><td>0</td><td>0.2</td></tr>\n",
              "    <tr><th>Dining Room</th><th>D</th><td>0</td><td>0</td><td>0</td><td>0</td><td>1</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n"
            ],
            "text/plain": [
              "<gtbook.display.pretty at 0x10cd0fd70>"
            ]
          },
          "execution_count": 67,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Great. Let's now represent \n",
        "\n",
        "room_states = ['Living Room', 'Kitchen', 'Office', 'Hallway', 'Dining Room']\n",
        "action_states = ['L', 'R', 'U', 'D'] # Left, Right, Up, Down\n",
        "# create a discrete variable with states\n",
        "\n",
        "#                   name,    states\n",
        "X = VARIABLES.discrete_series('X', range(1, 5), room_states)\n",
        "A = VARIABLES.discrete_series('A', range(1, 4), action_states)\n",
        "\n",
        "# conditional probability table for P(X2 | X1, A1)\n",
        "cond_prob = \"\"\"\n",
        "1/0/0/0/0 2/8/0/0/0 1/0/0/0/0 2/0/0/8/0\n",
        "8/2/0/0/0 0/1/0/0/0 0/1/0/0/0 0/2/0/0/8\n",
        "0/0/1/0/0 0/0/2/8/0 0/0/1/0/0 0/0/1/0/0\n",
        "0/0/8/2/0 0/0/0/2/8 8/0/0/2/0 0/0/0/1/0\n",
        "0/0/0/8/2 0/0/0/0/1 0/8/0/0/2 0/0/0/0/1\n",
        "\"\"\"\n",
        "\n",
        "conditional = gtsam.DiscreteConditional(X[2], [X[1], A[1]], cond_prob)\n",
        "\n",
        "# Add the conditional to the graph. Here we use pushback for adding an object with probabilities already specified. (In contrast to add which takes variable and probabilities/frequencies as input)\n",
        "graph = gtsam.DiscreteFactorGraph()\n",
        "graph.push_back(conditional)\n",
        "pretty(graph)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "It is also possible to represent the probabilities as a list/array to `gtsam.DiscreteConditional`. \n",
        "\n",
        "We just need to make sure to force cast the arguments be `key: tuple[int, int], parents: gtsam.gtsam.DiscreteKeys, table: list[float]`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 68,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div><p><tt>DiscreteFactorGraph</tt> of size 1</p><p>factor 0:</p><div>\n",
              "<p>  <i>P(X2|X1,A1):</i></p>\n",
              "<table class='DiscreteConditional'>\n",
              "  <thead>\n",
              "    <tr><th><i>X1</i></th><th><i>A1</i></th><th>Living Room</th><th>Kitchen</th><th>Office</th><th>Hallway</th><th>Dining Room</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>Living Room</th><th>L</th><td>1</td><td>8</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Living Room</th><th>R</th><td>0</td><td>2</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Living Room</th><th>U</th><td>0</td><td>0</td><td>1</td><td>8</td><td>0</td></tr>\n",
              "    <tr><th>Living Room</th><th>D</th><td>0</td><td>0</td><td>0</td><td>2</td><td>8</td></tr>\n",
              "    <tr><th>Kitchen</th><th>L</th><td>0</td><td>0</td><td>0</td><td>0</td><td>2</td></tr>\n",
              "    <tr><th>Kitchen</th><th>R</th><td>2</td><td>0</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Kitchen</th><th>U</th><td>8</td><td>1</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Kitchen</th><th>D</th><td>0</td><td>0</td><td>2</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Office</th><th>L</th><td>0</td><td>0</td><td>8</td><td>2</td><td>0</td></tr>\n",
              "    <tr><th>Office</th><th>R</th><td>0</td><td>0</td><td>0</td><td>8</td><td>1</td></tr>\n",
              "    <tr><th>Office</th><th>U</th><td>1</td><td>0</td><td>0</td><td>8</td><td>0</td></tr>\n",
              "    <tr><th>Office</th><th>D</th><td>0</td><td>1</td><td>0</td><td>0</td><td>8</td></tr>\n",
              "    <tr><th>Hallway</th><th>L</th><td>0</td><td>0</td><td>1</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Hallway</th><th>R</th><td>0</td><td>0</td><td>0</td><td>2</td><td>0</td></tr>\n",
              "    <tr><th>Hallway</th><th>U</th><td>0</td><td>0</td><td>0</td><td>0</td><td>2</td></tr>\n",
              "    <tr><th>Hallway</th><th>D</th><td>2</td><td>0</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Dining Room</th><th>L</th><td>0</td><td>2</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Dining Room</th><th>R</th><td>0</td><td>0</td><td>1</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Dining Room</th><th>U</th><td>8</td><td>0</td><td>0</td><td>1</td><td>0</td></tr>\n",
              "    <tr><th>Dining Room</th><th>D</th><td>0</td><td>8</td><td>0</td><td>0</td><td>1</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n"
            ],
            "text/plain": [
              "<gtbook.display.pretty at 0x10ce3cdd0>"
            ]
          },
          "execution_count": 68,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "cpt = [\n",
        "    [1, 0, 0, 0, 0], [2, 8, 0, 0, 0], [1, 0, 0, 0, 0], [2, 0, 0, 8, 0],\n",
        "    [8, 2, 0, 0, 0], [0, 1, 0, 0, 0], [0, 1, 0, 0, 0], [0, 2, 0, 0, 8],\n",
        "    [0, 0, 1, 0, 0], [0, 0, 2, 8, 0], [0, 0, 1, 0, 0], [0, 0, 1, 0, 0],\n",
        "    [0, 0, 8, 2, 0], [0, 0, 0, 2, 8], [8, 0, 0, 2, 0], [0, 0, 0, 1, 0],\n",
        "    [0, 0, 0, 8, 2], [0, 0, 0, 0, 1], [0, 8, 0, 0, 2], [0, 0, 0, 0, 1],\n",
        "]\n",
        "\n",
        "# flatten cpt and create motion model\n",
        "flatten_cpt = [item for sublist in cpt for item in sublist]\n",
        "\n",
        "# force case parents \n",
        "parents = gtsam.DiscreteKeys()\n",
        "parents.push_back(X[1])\n",
        "parents.push_back(A[1])\n",
        "\n",
        "motion_model2 = gtsam.DiscreteConditional(key=X[2], parents=parents, table=flatten_cpt)\n",
        "graph2 = gtsam.DiscreteFactorGraph()\n",
        "graph2.push_back(motion_model2)\n",
        "pretty(graph2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's look at how the graph looks like"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 69,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 12.0.0 (0)\n",
              " -->\n",
              "<!-- Pages: 1 -->\n",
              "<svg width=\"206pt\" height=\"84pt\"\n",
              " viewBox=\"0.00 0.00 206.00 83.60\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 79.6)\">\n",
              "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-79.6 202,-79.6 202,4 -4,4\"/>\n",
              "<!-- var4683743612465315841 -->\n",
              "<g id=\"node1\" class=\"node\">\n",
              "<title>var4683743612465315841</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-57.6\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-52.55\" font-family=\"Times,serif\" font-size=\"14.00\">A1</text>\n",
              "</g>\n",
              "<!-- factor0 -->\n",
              "<g id=\"node4\" class=\"node\">\n",
              "<title>factor0</title>\n",
              "<ellipse fill=\"black\" stroke=\"black\" cx=\"99\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
              "</g>\n",
              "<!-- var4683743612465315841&#45;&#45;factor0 -->\n",
              "<g id=\"edge3\" class=\"edge\">\n",
              "<title>var4683743612465315841&#45;&#45;factor0</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M44.43,-43.58C63.74,-29.15 92.87,-7.38 98.16,-3.43\"/>\n",
              "</g>\n",
              "<!-- var6341068275337658369 -->\n",
              "<g id=\"node2\" class=\"node\">\n",
              "<title>var6341068275337658369</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"99\" cy=\"-57.6\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"99\" y=\"-52.55\" font-family=\"Times,serif\" font-size=\"14.00\">X1</text>\n",
              "</g>\n",
              "<!-- var6341068275337658369&#45;&#45;factor0 -->\n",
              "<g id=\"edge2\" class=\"edge\">\n",
              "<title>var6341068275337658369&#45;&#45;factor0</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M99,-39.28C99,-25.77 99,-8.54 99,-3.96\"/>\n",
              "</g>\n",
              "<!-- var6341068275337658370 -->\n",
              "<g id=\"node3\" class=\"node\">\n",
              "<title>var6341068275337658370</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"171\" cy=\"-57.6\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"171\" y=\"-52.55\" font-family=\"Times,serif\" font-size=\"14.00\">X2</text>\n",
              "</g>\n",
              "<!-- var6341068275337658370&#45;&#45;factor0 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>var6341068275337658370&#45;&#45;factor0</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M153.57,-43.58C134.26,-29.15 105.13,-7.38 99.84,-3.43\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<gtbook.display.show at 0x10ce3ee10>"
            ]
          },
          "execution_count": 69,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "display.show(graph)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Sampling from Probability Distribution\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "If we just want a sample with any prior knowledge about our robot's location or actions, simply call `.sample` on the `gtsam.DiscreteFactorGraph` graph. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 70,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 12.0.0 (0)\n",
              " -->\n",
              "<!-- Pages: 1 -->\n",
              "<svg width=\"134pt\" height=\"116pt\"\n",
              " viewBox=\"0.00 0.00 134.00 116.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 112)\">\n",
              "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-112 130,-112 130,4 -4,4\"/>\n",
              "<!-- var4683743612465315841 -->\n",
              "<g id=\"node1\" class=\"node\">\n",
              "<title>var4683743612465315841</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">A1</text>\n",
              "</g>\n",
              "<!-- var6341068275337658369 -->\n",
              "<g id=\"node2\" class=\"node\">\n",
              "<title>var6341068275337658369</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-90\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">X1</text>\n",
              "</g>\n",
              "<!-- var6341068275337658369&#45;&gt;var4683743612465315841 -->\n",
              "<g id=\"edge2\" class=\"edge\">\n",
              "<title>var6341068275337658369&#45;&gt;var4683743612465315841</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M27,-71.83C27,-64.55 27,-55.98 27,-47.83\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"30.5,-47.93 27,-37.93 23.5,-47.93 30.5,-47.93\"/>\n",
              "</g>\n",
              "<!-- var6341068275337658370 -->\n",
              "<g id=\"node3\" class=\"node\">\n",
              "<title>var6341068275337658370</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"99\" cy=\"-90\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"99\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">X2</text>\n",
              "</g>\n",
              "<!-- var6341068275337658370&#45;&gt;var4683743612465315841 -->\n",
              "<g id=\"edge3\" class=\"edge\">\n",
              "<title>var6341068275337658370&#45;&gt;var4683743612465315841</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M83.73,-74.73C73.99,-64.99 61.19,-52.19 50.21,-41.21\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"52.86,-38.91 43.32,-34.32 47.91,-43.86 52.86,-38.91\"/>\n",
              "</g>\n",
              "<!-- var6341068275337658370&#45;&gt;var6341068275337658369 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>var6341068275337658370&#45;&gt;var6341068275337658369</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M71.78,-90C69.72,-90 67.62,-90 65.5,-90\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"65.77,-86.5 55.77,-90 65.77,-93.5 65.77,-86.5\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<gtbook.display.show at 0x10ce3cbc0>"
            ]
          },
          "execution_count": 70,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "posterior = graph.sumProduct() # Powerful fancy optimization happening here read more at https://www.roboticsbook.org/S34_vacuum_perception.html#sum-product-in-gtsam\n",
        "display.show(posterior, hints={\"X\": 1, \"A\":0})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 71,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<table class='DiscreteValues'>\n",
              "  <thead>\n",
              "    <tr><th>Variable</th><th>value</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>A1</th><td>U</td></tr>\n",
              "    <tr><th>X1</th><td>Office</td></tr>\n",
              "    <tr><th>X2</th><td>Office</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "<gtbook.display.pretty at 0x10a092ba0>"
            ]
          },
          "execution_count": 71,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "sample = posterior.sample()\n",
        "pretty(sample)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "If we do have some additional knowledge. Say we know that: \n",
        "1. robot started in the Living Room with 80% probability, 5% in the rest. So X[1] = [0.8, 0.05, 0.05, 0.05, 0.05]\n",
        "2. robot took action 'right' at time 1\n",
        "\n",
        "We can condition on that knowledge and sample again. To do this, we do have to build the graph again to incorporate these knowledge. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 72,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<table class='DiscreteValues'>\n",
              "  <thead>\n",
              "    <tr><th>Variable</th><th>value</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>X1</th><td>Kitchen</td></tr>\n",
              "    <tr><th>X2</th><td>Kitchen</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "<gtbook.display.pretty at 0x10ce2ac60>"
            ]
          },
          "execution_count": 72,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "graph = gtsam.DiscreteFactorGraph()\n",
        "\n",
        "# add information about X1\n",
        "graph.add(X[1], \"80 5 5 5 5\")  # P(X1)\n",
        "# add assignment of A1\n",
        "actions = VARIABLES.assignment({A[1]: 'R'}) # action A1 = 'R' \n",
        "conditional = gtsam.DiscreteConditional(X[2], [X[1], A[1]], cond_prob)\n",
        "conditional_up = conditional.choose(actions)  # new_conditional probabilities given that A1 = 'R'\n",
        "\n",
        "graph.push_back(conditional_up)\n",
        "posterior = graph.sumProduct() # Powerful fancy optimization happening here read more at https://www.roboticsbook.org/S34_vacuum_perception.html#sum-product-in-gtsam\n",
        "# show(posterior, hints={\"X\": 1})\n",
        "sample = posterior.sample()\n",
        "pretty(sample)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Notice how the variable A[1] is no longer present in our sample. (Since it's always fixed). If we examine `conditional` and `conditional_up`, we can see that we have selected only the rows with A[1] = 'R' and completely removed the column A[1]. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 73,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<p>  <i>P(X2|X1,A1):</i></p>\n",
              "<table class='DiscreteConditional'>\n",
              "  <thead>\n",
              "    <tr><th><i>X1</i></th><th><i>A1</i></th><th>Living Room</th><th>Kitchen</th><th>Office</th><th>Hallway</th><th>Dining Room</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>Living Room</th><th>L</th><td>1</td><td>0</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Living Room</th><th>R</th><td>0.2</td><td>0.8</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Living Room</th><th>U</th><td>1</td><td>0</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Living Room</th><th>D</th><td>0.2</td><td>0</td><td>0</td><td>0.8</td><td>0</td></tr>\n",
              "    <tr><th>Kitchen</th><th>L</th><td>0.8</td><td>0.2</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Kitchen</th><th>R</th><td>0</td><td>1</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Kitchen</th><th>U</th><td>0</td><td>1</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Kitchen</th><th>D</th><td>0</td><td>0.2</td><td>0</td><td>0</td><td>0.8</td></tr>\n",
              "    <tr><th>Office</th><th>L</th><td>0</td><td>0</td><td>1</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Office</th><th>R</th><td>0</td><td>0</td><td>0.2</td><td>0.8</td><td>0</td></tr>\n",
              "    <tr><th>Office</th><th>U</th><td>0</td><td>0</td><td>1</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Office</th><th>D</th><td>0</td><td>0</td><td>1</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Hallway</th><th>L</th><td>0</td><td>0</td><td>0.8</td><td>0.2</td><td>0</td></tr>\n",
              "    <tr><th>Hallway</th><th>R</th><td>0</td><td>0</td><td>0</td><td>0.2</td><td>0.8</td></tr>\n",
              "    <tr><th>Hallway</th><th>U</th><td>0.8</td><td>0</td><td>0</td><td>0.2</td><td>0</td></tr>\n",
              "    <tr><th>Hallway</th><th>D</th><td>0</td><td>0</td><td>0</td><td>1</td><td>0</td></tr>\n",
              "    <tr><th>Dining Room</th><th>L</th><td>0</td><td>0</td><td>0</td><td>0.8</td><td>0.2</td></tr>\n",
              "    <tr><th>Dining Room</th><th>R</th><td>0</td><td>0</td><td>0</td><td>0</td><td>1</td></tr>\n",
              "    <tr><th>Dining Room</th><th>U</th><td>0</td><td>0.8</td><td>0</td><td>0</td><td>0.2</td></tr>\n",
              "    <tr><th>Dining Room</th><th>D</th><td>0</td><td>0</td><td>0</td><td>0</td><td>1</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "<gtbook.display.pretty at 0x10ce3c620>"
            ]
          },
          "execution_count": 73,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "pretty(conditional)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 74,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<p>  <i>P(X2|X1):</i></p>\n",
              "<table class='DiscreteConditional'>\n",
              "  <thead>\n",
              "    <tr><th><i>X1</i></th><th>Living Room</th><th>Kitchen</th><th>Office</th><th>Hallway</th><th>Dining Room</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>Living Room</th><td>0.2</td><td>0.8</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Kitchen</th><td>0</td><td>1</td><td>0</td><td>0</td><td>0</td></tr>\n",
              "    <tr><th>Office</th><td>0</td><td>0</td><td>0.2</td><td>0.8</td><td>0</td></tr>\n",
              "    <tr><th>Hallway</th><td>0</td><td>0</td><td>0</td><td>0.2</td><td>0.8</td></tr>\n",
              "    <tr><th>Dining Room</th><td>0</td><td>0</td><td>0</td><td>0</td><td>1</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "<gtbook.display.pretty at 0x10ce3c950>"
            ]
          },
          "execution_count": 74,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "pretty(conditional_up)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Congratulations! \n",
        "You have built your first discrete bayes net in gtsam, and performed efficient inference with it."
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "py312",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}
